# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import hashlib
import logging
import random
import sys

import kafka
from absl import app, flags
from tornado import gen
from tornado.ioloop import IOLoop

from coin.util.queue.config import KafkaConfig


def fg(color):
  return '\033[38;5;%dm' % color


class KafkaMessagePrinter(object):
  def __init__(self, *, kafka_config, topics=None, color=True, randomize_color=False, logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._topics = kafka_config.gen_topics(topics)
    self._kafka_config = kafka_config
    assert kafka_config is not None
    self._kafka_consumer = None

    self._color = color
    self._color_offset = random.randint(0, 255) if randomize_color else 0

    self._observed_max_topic_len = 0

  def _start(self):
    # TODO(inkyu): Find offset by time
    self._kafka_consumer = kafka.KafkaConsumer(bootstrap_servers=self._kafka_config.kafka_servers,
                                               auto_offset_reset='latest')

    if self._topics is None:
      self._topics = self._kafka_consumer.topics()

    if not self._topics:
      raise KeyError('No topic')

    self._logger.info('Topics: %s', ', '.join(self._topics))
    self._kafka_consumer.subscribe(self._topics)

  @gen.coroutine
  def _loop(self):
    while True:
      msg = next(self._kafka_consumer)
      self._print_msg(msg)

  def start(self, *, ioloop=None):
    self._start()

    ioloop = ioloop or IOLoop.current()
    ioloop.add_callback(self._loop)

  def _topic_to_color_str(self, topic):
    hashed_topic = hashlib.md5(bytes(topic, 'ascii')).digest()
    color_no = (hashed_topic[4] + self._color_offset) % 256
    if color_no in (0, 15, 16, 231, 232, 233, 234, 235, 254, 255):
      color_no = (color_no + 87) % 256
    return fg(color_no)

  def _print_msg(self, msg):
    msg_dt = datetime.datetime.fromtimestamp(msg.timestamp / 1000.)
    try:
      msg_value = msg.value.decode('utf-8').rstrip()
    except UnicodeDecodeError:
      msg_value = str(msg.value)
    msg_lines = msg_value.split('\n')

    self._observed_max_topic_len = max(self._observed_max_topic_len, len(msg.topic))
    topic_len = min(80, self._observed_max_topic_len)
    msg_time_str = msg_dt.strftime('%m%d %H:%M:%S.%f')[:-3]

    if self._color:
      color_str = self._topic_to_color_str(msg.topic)
      format_str = '%s%%-%ds%s │ %s%%17s%s ┃ %s%%s%s' % (
          color_str, topic_len, '\033[0m', color_str, '\033[0m', color_str, '\033[0m')
    else:
      format_str = '%%-%ds │ %%17s ┃ %%s' % topic_len

    for lineno, line in enumerate(msg_lines):
      if lineno == 0:
        print(format_str % (msg.topic[:topic_len], msg_time_str, line))
      else:
        print(format_str % ('', '', line))


def main(argv):
  FLAGS = flags.FLAGS
  kafka_config = KafkaConfig.from_cmd_config(FLAGS.kafka_config_filename)
  topics = FLAGS.topic.split(',') if FLAGS.topic else None

  ioloop = IOLoop()
  msg_printer = KafkaMessagePrinter(kafka_config=kafka_config,
                                    topics=topics,
                                    color=(sys.stdout.isatty() and FLAGS.color),
                                    randomize_color=FLAGS.randomize_color)
  msg_printer.start(ioloop=ioloop)

  try:
    ioloop.start()
  except (KeyboardInterrupt, SystemExit):
    print()
    return


if __name__ == "__main__":
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')

  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')

  flags.DEFINE_string('topic', None, 'Topic to dump. All topics if None')

  flags.DEFINE_bool('color', True, 'If true, print message in color.')

  flags.DEFINE_bool('randomize_color',
                    False,
                    'If true, color corresponds to a topic will be changed each launch.')

  logging.getLogger('kafka').setLevel(logging.WARNING)
  app.run(main)
